import openai
import os


def parse_goal_direction(language_instr):
    """
    解析语言指令并返回目标方向
    
    Args:
        language_instr (str): 语言指令，描述如何到达目标位置
                            例如: "向前走到树那里，向右转，继续走直到找到一辆车"
    
    Returns:
        str: 目标方向，可能的值包括：
             - "front": 前方
             - "back": 后方  
             - "left": 左方
             - "right": 右方
             - "left front": 左前方
             - "right front": 右前方
             - "left back": 左后方
             - "right back": 右后方
    
    Raises:
        ValueError: 当OPENAI_KEY环境变量未设置时
        Exception: 当API调用失败时
    """
    # 检查环境变量
    # if "OPENAI_KEY" not in os.environ:
    #     raise ValueError("OPENAI_KEY环境变量未设置")
    
    #openai_key = os.environ["OPENAI_KEY"]
    openai_key = "EMPTY"
    openai.api_key = openai_key
    openai.base_url = "http://localhost:11434/v1/"
    
    try:
        response = openai.chat.completions.create(
            model="qwen2.5vl:7b",
            messages=[
                # 基本方向样本
                {
                    "role": "user",
                    "content": "Go straight ahead to the building, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "front"
                },
                {
                    "role": "user",
                    "content": "Turn around and walk to the entrance, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "back"
                },
                {
                    "role": "user",
                    "content": "Turn left and go to the store, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "left"
                },
                {
                    "role": "user",
                    "content": "Turn right and proceed to the parking lot, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "right"
                },
                
                # 复合方向样本
                {
                    "role": "user",
                    "content": "Walk forward to the tree, then turn left and continue to the bench, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "left front"
                },
                {
                    "role": "user",
                    "content": "Go straight to the intersection, turn right and walk to the restaurant, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "right front"
                },
                {
                    "role": "user",
                    "content": "Turn left, walk straight, then turn left again to reach the library, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "left back"
                },
                {
                    "role": "user",
                    "content": "Turn right, go forward, then turn right again to find the cafe, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "right back"
                },
                
                # 复杂路径样本
                {
                    "role": "user",
                    "content": "Move forward to the traffic light, turn left, walk until you see a bank, turn right and go to the post office, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "left front"
                },
                {
                    "role": "user",
                    "content": "Go straight to the park, turn left, walk until you find a school, turn left again to reach the museum, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "left back"
                },
                {
                    "role": "user",
                    "content": "Turn right, proceed to the mall, turn right again, walk until you see a theater, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "right back"
                },
                
                # 特殊情况样本
                {
                    "role": "user",
                    "content": "Head towards the mountain in front of you, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "front"
                },
                {
                    "role": "user",
                    "content": "Go back to where you came from, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "back"
                },
                {
                    "role": "user",
                    "content": "Take a left turn and head towards the lake, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "left"
                },
                {
                    "role": "user",
                    "content": "Make a right turn and go to the bridge, where is the destination"
                },
                {
                    "role": "assistant",
                    "content": "right"
                },
                
                # 当前查询
                {
                    "role": "user",
                    "content": f"{language_instr}, where is the destination"
                }
            ],
            max_tokens=300,
        )

        text = response.choices[0].message.content.strip()
        return text
        
    except Exception as e:
        raise Exception(f"解析方向失败: {str(e)}")

def validate_direction(direction):
    """
    验证返回的方向是否有效
    
    Args:
        direction (str): 要验证的方向字符串
    
    Returns:
        bool: 如果方向有效返回True，否则返回False
    """
    valid_directions = [
        "front", "back", "left", "right",
        "left front", "right front", "left back", "right back"
    ]
    return direction.lower() in valid_directions

def analyze_sample_coverage():
    """
    分析few-shot样本的覆盖情况
    
    Returns:
        dict: 包含各方向的样本数量和覆盖情况
    """
    # 从函数中提取样本（这里手动统计）
    samples = {
        "front": 2,      # "Go straight ahead to the building", "Head towards the mountain"
        "back": 2,       # "Turn around and walk to the entrance", "Go back to where you came from"
        "left": 2,       # "Turn left and go to the store", "Take a left turn and head towards the lake"
        "right": 2,      # "Turn right and proceed to the parking lot", "Make a right turn and go to the bridge"
        "left front": 3, # 复合方向样本
        "right front": 3,
        "left back": 2,
        "right back": 2
    }
    
    total_samples = sum(samples.values())
    coverage = {direction: count for direction, count in samples.items()}
    coverage["total"] = total_samples
    
    return coverage

if __name__ == "__main__":
    # 显示样本覆盖情况
    print("=== Few-shot样本覆盖情况 ===")
    coverage = analyze_sample_coverage()
    for direction, count in coverage.items():
        if direction != "total":
            print(f"{direction}: {count}个样本")
    print(f"总计: {coverage['total']}个样本")
    print("=" * 50)
    
    # 测试代码
    test_instructions = [
        # 基本方向测试
        "Walk straight to the library, turn left, and proceed until you reach a bus stop",
        "Turn right and go to the store",
        "Go straight ahead to the building",
        "Turn around and walk back to the entrance",
        
        # 复合方向测试
        "Go forward to the intersection, turn left and continue to the park",
        "Turn right, walk straight, then turn right again to find the cafe",
        "Walk forward to the tree, then turn left and continue to the bench",
        "Turn left, walk straight, then turn left again to reach the museum",
        
        # 复杂路径测试
        "Move forward to the traffic light, turn left, walk until you see a bank, turn right and go to the post office",
        "Turn right at the corner, walk straight, turn left at the next intersection, continue until you reach the hospital",
    
    ]
    
    print("\n=== 测试结果 ===")
    for i, instruction in enumerate(test_instructions, 1):
        try:
            result = parse_goal_direction(instruction)
            print(f"测试 {i}: {instruction}")
            print(f"结果: {result}")
            print(f"有效: {validate_direction(result)}")
            print("-" * 50)
        except Exception as e:
            print(f"测试 {i}: 处理指令 '{instruction}' 时出错: {e}")
            print("-" * 50)
